﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using TorchSharp;
using static TorchSharp.torch.nn;
using static TorchSharp.torch;
using TorchSharp.Modules;

namespace NgpSharp
{
    public class SHEncoder : torch.nn.Module<torch.Tensor, int, torch.Tensor>
    {
        private int input_dim;
        private int degree;
        public int output_dim {  get; private set; }

        public SHEncoder(string name, int input_dim = 3, int degree = 4) 
            : base(name)
        {
            this.input_dim = input_dim;
            this.degree = degree;
            output_dim = (int)Math.Pow(degree, 2);

            if (input_dim != 3) throw new ArgumentException("SH encoder only support input dim == 3");
            if (!(degree > 0 && degree <= 5)) throw new ArgumentException("SH encoder only supports degree in [1, 5]");
            RegisterComponents();
        }

        public override torch.Tensor forward(torch.Tensor input, int size)
        {
            using var c = torch.NewDisposeScope();

            input = input / size;
            List<long> prefix_shape = [..input.shape[..^1]];
            input = input.reshape(-1, input_dim);

            Tensor outputs = SHEncoderFunction.apply(input, (uint)degree, false);
            prefix_shape.Add(output_dim);
            outputs = outputs.reshape([.. prefix_shape]);
            return outputs.MoveToOuterDisposeScope();
        }
    }
}
